import traceback
from datetime import datetime

import pandas as pd
from xbase_util.common_util import date2s
from xbase_util.db.bean import ConfigBean

from base.bean.beans import ProcedureEnum
from base.bean.status_enum import StatusEnum
from base.util.redis_util import UpdateStatusParams


def sample_fields(df1, df2, mapping_dict):
    map_list = mapping_dict.split(";")
    df1_extended = df1.sample(n=len(df2), replace=True, random_state=42).reset_index(drop=True)
    for df_value in map_list:
        df_v_list = df_value.split(":")
        df2_col = df_v_list[0]
        df1_col = df_v_list[1]
        if df1_col in df1.columns and df2_col in df2.columns:
            df2[df2_col] = df1_extended[df1_col]
    return df2


def replace_fields(config: ConfigBean, queue, geoUtil) -> bool:
    queue.put(UpdateStatusParams(status=StatusEnum.processing.value, procedure=ProcedureEnum.replace_fields.value,
                                 process=0))
    try:
        destination_file = pd.read_csv(f"{config.replace_destination}")
        final_file = pd.DataFrame()
        source = config.replace_source.split(',')
        for i, path in enumerate(source):
            print(f"[替换五元组]正在读取文件{i + 1}/{len(source)}|{path}")
            file = pd.read_csv(f"{path}")
            final_file = pd.concat([file, final_file], axis=0, ignore_index=True)
        print("[替换五元组]开始")
        df = sample_fields(final_file, destination_file, config.replace_mapping)
        df = df.apply(lambda x: geoUtil.get_geo_by_ip(x), axis=1)
        df.to_csv(f"{config.replace_save_to}", index=False)
        print("[替换五元组]替换完毕")
        return True
    except Exception as e:
        print(e)
        traceback.print_exc()
        queue.put(UpdateStatusParams(status=StatusEnum.error.value, error=f"{e}",
                                     procedure=ProcedureEnum.replace_fields.value,
                                     execEndTime=date2s(datetime.now())))
        queue.put(None)
        return False
